package site.wanjiahao.tree;

/**
 * 红黑树(自平衡二叉搜索树)，类似于4阶B树
 */
public class RBTree<E> extends BST<E> {

    private static final int BLACK = 1;

    private static final int RED = 0;

    @Override
    protected void addAfter(Node node) {
        // 如果是根节点
        Node parent = node.parent;
        if (parent == null) {
            // 染色
            black(node);
            return;
        }

        // 如果父节点为黑色
        if (isBlack(parent)) {
            // 直接添加，不做处理
            return;
        }

        // 父节点是红色
        Node uncle = parent.sibling();
        // 祖父染红
        Node grand = red(parent.parent);
        if (isBlack(uncle)) {
            // uncle 不是红色 旋转处理
            if (parent.isLeftNode()) {
                if (node.isRightNode()) {
                    // LR
                    parent.rotateLeft();
                    parent = node;
                }
                // LL
                grand.rotateRight();
            } else {
                if (node.isLeftNode()) {
                    // RL
                    parent.rotateRight();
                    parent = node;
                }
                // RR
                grand.rotateLeft();
            }
            // 父亲染黑
            black(parent);
        } else {
            // uncle 是红色 上溢处理
            black(parent);
            black(uncle);
            addAfter(grand);
        }
    }

    @Override
    protected void removeAfter(Node node, Node replacement) {
        Node parent = node.parent;

        if (isRed(node)) {
            // 删除红色节点，不用做任何处理
            return;
        }

        if (replacement != null) {
            // 删除黑色节点，并且黑色节点存在替代节点
            // 替代节点染黑，其余不用做任何处理
            black(replacement);
            return;
        }

        // 删除的是根节点
        if (parent == null) return;

        // 删除单个黑色节点，做下移处理
        boolean left = parent.left == null || node.isLeftNode();
        Node sibling = left ? parent.right : parent.left;
        if (left) {
            // 当前节点在左边
            if (isRed(sibling)) {
                // 兄弟节点是红色
                black(sibling);
                red(parent);
                // 变成兄弟节点为黑节点情况
                parent.rotateLeft();
                sibling = parent.right;
            }
            if (sibling == null) return;
            // 兄弟节点是黑节点
            if (isRed(sibling.left) || isRed(sibling.right)) {
                if (isBlack(sibling.right)) {
                    // RL
                    sibling.rotateRight();
                    sibling = parent.right;
                }
                // RR
                parent.rotateLeft();
                // 染色
                color(sibling, colorOf(parent));
                color(sibling.left, RED);
                color(sibling.right, RED);
            } else {
                // 左右无节点，合并操作
                if (isBlack(parent)) {
                    red(sibling);
                    removeAfter(parent, null);
                } else {
                    black(parent);
                    red(sibling);
                }
            }
        } else {
            // 当前节点在右边
            if (isRed(sibling)) {
                // 兄弟节点是红色
                black(sibling);
                red(parent);
                // 变成兄弟节点为黑节点情况
                parent.rotateRight();
                sibling = parent.left;
            }
            if (sibling == null) return;
            // 兄弟节点是黑节点
            if (isRed(sibling.left) || isRed(sibling.right)) {
                if (isBlack(sibling.left)) {
                    // LR
                    sibling.rotateLeft();
                    sibling = parent.left;
                }
                // LL
                parent.rotateRight();
                // 染色
                color(sibling, colorOf(parent));
                color(sibling.left, RED);
                color(sibling.right, RED);
            } else {
                // 左右无节点，合并操作
                if (isBlack(parent)) {
                    red(sibling);
                    removeAfter(parent, null);
                } else {
                    black(parent);
                    red(sibling);
                }
            }
        }
    }

    // 自定义工具方法
    private int colorOf(Node node) {
        return node == null ? BLACK : toRBNode(node).color;
    }

    private RBNode color(Node node, int color) {
        RBNode erbNode = toRBNode(node);
        if (node != null) {
            erbNode.color = color;
        }
        return erbNode;
    }

    private boolean isBlack(Node node) {
        return colorOf(node) == BLACK;
    }

    private boolean isRed(Node node) {
        return colorOf(node) == RED;
    }

    private RBNode black(Node node) {
        return color(node, BLACK);
    }

    private RBNode red(Node node) {
        return color(node, RED);
    }

    private RBNode toRBNode(Node node) {
        return ((RBNode) node);
    }


    @Override
    protected Node createNode(E ele, Node parent) {
        return new RBNode(ele, parent);
    }

    private class RBNode extends Node {

        // 默认添加红色节点，比较符合红黑树性质
        int color = RED;

        public RBNode(E ele, Node parent) {
            super(ele, parent);
        }

        @Override
        public String toString() {
            return ele + "_" + (color == RED ? "red" : "");
        }
    }
}
